from mindspore.dataset import vision, transforms
from mindspore.dataset.vision import Inter
from mindspore.dataset.vision import py_transforms_util as P
import numpy as np
import numbers
import random
from PIL import ImageFilter, ImageOps


class ColorJitter(object):
    """Randomly change the brightness, contrast, saturation and hue of an image.
    If the image is torch Tensor, it is expected
    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, mode "1", "L", "I", "F" and modes with transparency (alpha channel) are not supported.

    Args:
        brightness (float or tuple of float (min, max)): How much to jitter brightness.
            brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
            or the given [min, max]. Should be non negative numbers.
        contrast (float or tuple of float (min, max)): How much to jitter contrast.
            contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]
            or the given [min, max]. Should be non negative numbers.
        saturation (float or tuple of float (min, max)): How much to jitter saturation.
            saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]
            or the given [min, max]. Should be non negative numbers.
        hue (float or tuple of float (min, max)): How much to jitter hue.
            hue_factor is chosen uniformly from [-hue, hue] or the given [min, max].
            Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5.
    """

    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
        super().__init__()
        self.brightness = self._check_input(brightness, 'brightness')
        self.contrast = self._check_input(contrast, 'contrast')
        self.saturation = self._check_input(saturation, 'saturation')
        self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5),
                                     clip_first_on_zero=False)

    @staticmethod
    def _check_input(value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True):
        if isinstance(value, numbers.Number):
            if value < 0:
                raise ValueError("If {} is a single number, it must be non negative.".format(name))
            value = [center - float(value), center + float(value)]
            if clip_first_on_zero:
                value[0] = max(value[0], 0.0)
        elif isinstance(value, (tuple, list)) and len(value) == 2:
            if not bound[0] <= value[0] <= value[1] <= bound[1]:
                raise ValueError("{} values should be between {}".format(name, bound))
        else:
            raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name))

        # if value is 0 or (1., 1.) for brightness/contrast/saturation
        # or (0., 0.) for hue, do nothing
        if value[0] == value[1] == center:
            value = None
        return value

    @staticmethod
    def get_params(brightness, contrast, saturation, hue):
        """Get the parameters for the randomized transform to be applied on image.

        Args:
            brightness (tuple of float (min, max), optional): The range from which the brightness_factor is chosen
                uniformly. Pass None to turn off the transformation.
            contrast (tuple of float (min, max), optional): The range from which the contrast_factor is chosen
                uniformly. Pass None to turn off the transformation.
            saturation (tuple of float (min, max), optional): The range from which the saturation_factor is chosen
                uniformly. Pass None to turn off the transformation.
            hue (tuple of float (min, max), optional): The range from which the hue_factor is chosen uniformly.
                Pass None to turn off the transformation.

        Returns:
            tuple: The parameters used to apply the randomized transform
            along with their random order.
        """
        fn_idx = np.random.permutation(4)

        b = None if brightness is None else float(np.random.uniform(brightness[0], brightness[1]))
        c = None if contrast is None else float(np.random.uniform(contrast[0], contrast[1]))
        s = None if saturation is None else float(np.random.uniform(saturation[0], saturation[1]))
        h = None if hue is None else float(np.random.uniform(hue[0], hue[1]))

        return fn_idx, b, c, s, h

    def __call__(self, img):
        """
        Args:
            img (PIL Image or Tensor): Input image.

        Returns:
            PIL Image or Tensor: Color jittered image.
        """
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

        for fn_id in fn_idx:
            if fn_id == 0 and brightness_factor is not None:
                img = P.adjust_brightness(img, brightness_factor)
            elif fn_id == 1 and contrast_factor is not None:
                img = P.adjust_contrast(img, contrast_factor)
            elif fn_id == 2 and saturation_factor is not None:
                img = P.adjust_saturation(img, saturation_factor)
            elif fn_id == 3 and hue_factor is not None:
                img = P.adjust_hue(img, hue_factor)

        return img


def convert_RGB(image):
    return image.convert("RGB")


class GaussianBlur(object):
    """
    Apply Gaussian Blur to the PIL image.
    """

    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
        self.prob = p
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        do_it = random.random() <= self.prob
        if not do_it:
            return img

        return img.filter(
            ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)
            )
        )


class Solarization(object):
    """
    Apply Solarization to the PIL image.
    """

    def __init__(self, p):
        self.p = p

    def __call__(self, img):
        if random.random() < self.p:
            return ImageOps.solarize(img)
        else:
            return img


def copy(x):
    return x


class DataAugmentationDINO_Cluster(object):
    def __init__(self, global_crops_scale, local_crops_scale, local_crops_number):
        self.clip_transform = transforms.Compose([
            vision.Resize(224, interpolation=Inter.BICUBIC),  # Image.BICUBIC
            vision.CenterCrop(224),
            convert_RGB,
            vision.Normalize(
                (0.48145466, 0.4578275, 0.40821073),
                (0.26862954, 0.26130258, 0.27577711),
            ),
            vision.ToTensor(),
        ]
        )
        # =========cluster transform=============
        self.img2_transform = transforms.Compose([
            vision.Resize(224, interpolation=Inter.BICUBIC),  # Image.BICUBIC
            vision.CenterCrop(224),
            # vision.Decode(),
            ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
            vision.Normalize(  # 加上normalize
                (0.48145466, 0.4578275, 0.40821073),
                (0.26862954, 0.26130258, 0.27577711),
            ),
            vision.ToTensor(),
        ]
        )
        flip_and_color_jitter = transforms.Compose([  ##ColorJitter重写
            vision.RandomHorizontalFlip(prob=0.5),
            ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1),
            vision.RandomGrayscale(prob=0.2),
        ])
        normalize = transforms.Compose([
            vision.Normalize((0.48145466, 0.4578275, 0.40821073),
                             (0.26862954, 0.26130258, 0.27577711)),
            vision.ToTensor(),
        ])

        # first global crop
        self.global_transform1 = transforms.Compose([
            convert_RGB,
            vision.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Inter.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            normalize,
        ])
        # second global crop
        self.global_transform2 = transforms.Compose([
            convert_RGB,
            vision.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Inter.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(0.1),
            Solarization(0.2),
            normalize,
        ])
        # transformation for the local small crops
        self.local_crops_number = local_crops_number
        self.local_transform = transforms.Compose([
            convert_RGB,
            vision.RandomResizedCrop(96, scale=local_crops_scale, interpolation=Inter.BICUBIC),
            flip_and_color_jitter,
            GaussianBlur(p=0.5),
            normalize,
        ])

    def __call__(self, image, token):
        crops1 = []
        crops2 = []
        image = vision.Decode(True)(image)
        crops1.append(self.global_transform1(image)[0])
        crops1.append(self.global_transform2(image)[0])
        for _ in range(self.local_crops_number):
            crops2.append(self.local_transform(image)[0])
        crops1.append(self.clip_transform(image)[0])

        # =============img2 and affine2_to_1=================
        img2 = self.img2_transform(image)[0]
        # affine_kwargs = {"min_rot": -30, "max_rot": 30,
        #                  "min_shear": -10,
        #                  "max_shear": 10,
        #                  "min_scale": 0.8,
        #                  "max_scale": 1.2}
        # img2, affine1_to_2, affine2_to_1 = random_affine(img2[0], **affine_kwargs)
        # if np.random.rand() > 0.5:
        #     img2 = self.flip(img2)
        #     affine2_to_1[0, :] *= -1.
        crops1.append(img2)
        # crops.append(affine2_to_1)
        return crops1, crops2, token
